# %matplotlib inline
import os
import datetime

import torch
# from torch import nn
# from torch.nn import functional as F
# import torchvision
# from d2l import torch as d2l
# import pandas as pd

from ssd import TinySSD
from dataset import BananasDataset, train_transforms
from loss import multibox_target, calc_loss
from eval import cls_eval, bbox_eval, Accumulator
from visual import Animator, Timer
from hardware import try_gpu

os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"

# Y1 = forward(torch.zeros((2, 8, 20, 20)), cls_predictor(8, 5, 10))
# Y2 = forward(torch.zeros((2, 16, 10, 10)), cls_predictor(16, 3, 10))
# Y1.shape, Y2.shape

# forward(torch.zeros((2, 3, 256, 256)), base_net()).shape

net = TinySSD(num_classes=1)
X = torch.zeros((32, 3, 256, 256))
anchors, cls_preds, bbox_preds = net(X)

print('output anchors:', anchors.shape)
print('output class preds:', cls_preds.shape)
print('output bbox preds:', bbox_preds.shape)

# 训练模型
batch_size = 32
# train_iter, _ = d2l.load_data_bananas(batch_size)
train_iter = torch.utils.data.DataLoader(BananasDataset(is_train=True,
                                                        transforms=train_transforms),
                                         batch_size, shuffle=True)


# device = try_gpu()
device = 'cpu'

net = TinySSD(num_classes=1)
trainer = torch.optim.SGD(net.parameters(), lr=0.2, weight_decay=5e-4)

# num_epochs, timer = 20, d2l.Timer()
num_epochs, timer = 20, Timer()

# 动画器
# animator = d2l.Animator(xlabel='epoch', xlim=[1, num_epochs],
# animator = Animator(xlabel='epoch', xlim=[1, num_epochs],
#                     legend=['class error', 'bbox mae'])
net = net.to(device)

for epoch in range(num_epochs):
    # 训练精确度的和，训练精确度的和中的示例数
    # 绝对误差的和，绝对误差的和中的示例数
    # metric = d2l.Accumulator(4)
    metric = Accumulator(4)  # 累加器
    net.train()
    for features, target in train_iter:
        timer.start()
        trainer.zero_grad()
        X, Y = features.to(device), target.to(device)
        # 生成多尺度的锚框，为每个锚框预测类别和偏移量
        anchors, cls_preds, bbox_preds = net(X)
        # 为每个锚框标注类别和偏移量
        # bbox_labels, bbox_masks, cls_labels = d2l.multibox_target(anchors, Y)
        bbox_labels, bbox_masks, cls_labels = multibox_target(anchors, Y)
        # 根据类别和偏移量的预测和标注值计算损失函数
        l = calc_loss(cls_preds, cls_labels, bbox_preds, bbox_labels, bbox_masks)
        l.mean().backward()
        trainer.step()
        metric.add(cls_eval(cls_preds, cls_labels), cls_labels.numel(),
                   bbox_eval(bbox_preds, bbox_labels, bbox_masks), bbox_labels.numel())
    cls_err, bbox_mae = 1 - metric[0] / metric[1], metric[2] / metric[3]
    # animator.add(epoch + 1, (cls_err, bbox_mae))
    print(f'epoch: {epoch+1}, class err: {cls_err:.2e}, bbox mae: {bbox_mae:.2e}')

torch.save(
    net.state_dict(), 
    'TinySSD'+datetime.datetime.now().strftime('%Y-%m-%d-%H_%M_%S')+'.pth'
)
print(f'{len(train_iter.dataset) / timer.stop():.1f} examples/sec on '
      f'{str(device)}')
